import * as lancedb from "@lancedb/lancedb";
import { env, pipeline } from "@xenova/transformers";

const databaseDir = "./db/lancedb";
const tableName = "codebase_v2"; // 使用新表名避免与旧表冲突

/**
 * 初始化embed模型
 * @returns
 */
async function initEmbedModel() {
  env.allowLocalModels = true;
  env.allowRemoteModels = false;
  env.localModelPath = "models";

  // 创建一个来自管道的嵌入函数，该函数从批次中返回一个向量列表。
  // sourceColumn 是要嵌入的数据中的列的名称。
  // 管道的输出是一个张量 { data: Float32Array(384) }，因此要过滤出向量。
  let pipe = await pipeline("feature-extraction", "Xenova/all-MiniLM-L6-v2");
  return {
    sourceColumn: "", // 嵌入的列名
    embed: async function (batch) {
      let result = [];
      for (let text of batch) {
        const res = await pipe(text, { pooling: "mean", normalize: true });
        result.push(res.tolist());
      }
      return result;
    },
  };
}

/**
 * 将文本内容转换为向量并保存到 LanceDB
 * @param {*} data []Object
 */
async function saveLancedb(data) {
  const embedModel = await initEmbedModel();
  const validData = [];
  
  for (const row of data) {
    // 确保文本不为null或undefined，优先使用text，如果没有则使用code
    const text = row.text || row.code || '';
    if (text) {
      try {
        // 计算向量
        const vectors = await embedModel.embed([text]);
        
        console.log('原始向量输出:', JSON.stringify(vectors).substring(0, 100) + '...');
        
        // 打印完整的向量结构用于调试
        console.log('向量完整结构:', JSON.stringify({
          type: Array.isArray(vectors) ? 'array' : typeof vectors,
          length: Array.isArray(vectors) ? vectors.length : undefined,
          firstItemType: Array.isArray(vectors) && vectors.length > 0 ? 
            (Array.isArray(vectors[0]) ? 'array' : typeof vectors[0]) : undefined,
          firstItemLength: Array.isArray(vectors) && vectors.length > 0 && Array.isArray(vectors[0]) ? 
            vectors[0].length : undefined
        }));
        
        // 简化的向量提取逻辑
        let validVector = null;
        
        // 处理可能的各种嵌套情况
        if (Array.isArray(vectors)) {
          // 情况1: vectors是二维数组，如[[vector]]
          if (vectors.length > 0 && Array.isArray(vectors[0])) {
            // 检查子数组是否是有效向量(长度>1)
            if (vectors[0].length > 1) {
              validVector = vectors[0];
            }
            // 情况2: 更复杂的嵌套，如[[[vector]]]
            else if (vectors[0].length > 0 && Array.isArray(vectors[0][0]) && vectors[0][0].length > 1) {
              validVector = vectors[0][0];
            }
          }
          // 情况3: vectors已经是一维向量数组
          else if (vectors.length > 1) {
            validVector = vectors;
          }
        }
        
        if (validVector && validVector.length > 1) {
          console.log(`成功提取有效向量，维度: ${validVector.length}`);
          validData.push({
            ...row,
            vector: validVector
          });
        } else {
          console.warn('无法提取有效向量，跳过此文本:', text.substring(0, 50) + '...');
        }
      } catch (error) {
        console.warn('向量计算出错:', error.message);
      }
    } else {
        // 详细记录空文本情况
        console.warn('跳过空文本的向量计算:', {
          hasText: !!row.text,
          hasCode: !!row.code,
          path: row.path || 'unknown'
        });
      }
  }
  
  console.log(`共处理 ${data.length} 个数据项，成功提取 ${validData.length} 个有效向量`);
  
  console.log(`准备保存 ${validData.length} 个有效向量到LanceDB`);
  if (validData.length > 0) {
    const db = await lancedb.connect(databaseDir);
    const table = await db.createTable(tableName, validData, { mode: "overwrite" });
    console.log('向量数据保存成功');
    return true; // 返回成功
  } else {
    console.log('没有有效向量数据可保存');
    console.warn('警告: 所有数据项都未能提取有效向量，可能是输入数据格式有误或全部为空文本');
    return false; // 返回失败
  }
}

/**
 * 检索LanceDB中的内容
 * @param query
 * @param topK
 */
async function searchLancedb(query, topK = 5) {
  const embedModel = await initEmbedModel();
  // 查询相似内容
  let queryEmbeddings = await embedModel.embed([query]); // 修复：确保传入数组
  let queryEmbed;
  
  // 处理嵌套的向量结构，与saveLancedb中的逻辑保持一致
  if (Array.isArray(queryEmbeddings)) {
    if (queryEmbeddings.length > 0 && Array.isArray(queryEmbeddings[0])) {
      if (queryEmbeddings[0].length > 1) {
        queryEmbed = queryEmbeddings[0];
      } else if (queryEmbeddings[0].length > 0 && Array.isArray(queryEmbeddings[0][0]) && queryEmbeddings[0][0].length > 1) {
        queryEmbed = queryEmbeddings[0][0];
      }
    } else if (queryEmbeddings.length > 1) {
      queryEmbed = queryEmbeddings;
    }
  }
  
  if (!queryEmbed || queryEmbed.length <= 1) {
    throw new Error('无法生成有效的查询向量');
  }
  
  const db = await lancedb.connect(databaseDir);
  const _tb = await db.openTable(tableName);
  const results = await _tb
    .vectorSearch(queryEmbed)
    .column("vector") // 明确指定向量列
    // .where("type == 'fruit'") // 可添加过滤条件
    .distanceType("cosine")
    .limit(topK)
    .toArray();

  return results.map((r) => r.text);
}

export { saveLancedb, searchLancedb };